import abc
from collections.abc import Iterable, Iterator
from typing import Any, TypeVar

from .iterator import DatasetIterator

AnySample = TypeVar('AnySample')


class Dataset(Iterable[AnySample], abc.ABC):
    """ A Dataset is a set of samples

    Definitely, it supports:
        - Iterate samples in order with an iterator.

    Optionally, it supports:
        - Getting a specific sample by index.
        - Estimating the total number of samples.
        - Dumping and resuming the state of its iterator.
    """

    def __len__(self) -> int:
        raise TypeError(f"{type(self)} does not support estimating the total number of samples.")

    def __getitem__(self, index: int, /) -> AnySample:
        raise TypeError(f"{type(self)} does not support getting a specific sample by index.")

    @abc.abstractmethod
    def __iter__(self) -> DatasetIterator[AnySample]:
        raise NotImplementedError

    def __resume__(self, state: Any | None = None, /) -> DatasetIterator[AnySample]:
        if state is None:
            return self.__iter__()
        raise TypeError(f"{self!r} does not support resuming!")


def resume(
    iterable: Dataset[AnySample] | Iterable[AnySample], state: Any | None = None, /
) -> DatasetIterator[AnySample] | Iterator[AnySample]:
    if isinstance(iterable, Dataset):
        return iterable.__resume__(state)
    if state is None:
        return iter(iterable)
    from zkl_aiutils_datasets.wrapping import wrap_dataset
    iterable = wrap_dataset(iterable)
    return resume(iterable, state)
